package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"time"

	"github.com/streadway/amqp"
)

type Target struct {
	IP   string
	Port string
}

type Config struct {
	Host     string
	Port     string
	Username string
	Password string
	Queue    string
}

func failOnError(err error, msg string) {
	if err != nil {
		log.Fatalf("%s: %s", msg, err)
	}
}

func scan(ip, port string) string {
	// 进行端口扫描
	conn, _ := net.DialTimeout(
		"tcp",
		fmt.Sprintf("%s:%s", ip, port),
		2*time.Second,
	)
	if conn != nil {
		conn.Close()
		return "Open"
	} else {
		return "Closed"
	}
}

func LoadConfig(path string, conf *Config) {
	buf, err := ioutil.ReadFile(path)
	failOnError(err, fmt.Sprintf("Failed to load %s", path))
	err = json.Unmarshal(buf, conf)
	failOnError(err, "Failed to read config")
}

func main() {
	var conf Config
	LoadConfig("conf.json", &conf)
	url := fmt.Sprintf("amqp://%s:%s@%s:%s/", conf.Username, conf.Password, conf.Host, conf.Port)
	conn, err := amqp.Dial(url)
	failOnError(err, "Failed to connect to RabbitMQ")
	defer conn.Close()
	ch, err := conn.Channel()
	failOnError(err, "Failed to open a channel")
	defer ch.Close()
	q, err := ch.QueueDeclare(
		conf.Queue, // name
		false,      // durable
		false,      // delete when unused
		false,      // exclusive
		false,      // no-wait
		nil,        // arguments
	)
	failOnError(err, "Failed to declare a queue")

	err = ch.Qos(
		100,   // perfetch count
		0,     // prefetch size
		false, // global
	)
	failOnError(err, "Failed to set Qos")
	msgs, err := ch.Consume(
		q.Name, // queue
		"",     // consumer
		false,  // auto-ack
		false,  // exclusive
		false,  // no-local
		false,  // no-wait
		nil,    // args
	)
	failOnError(err, "Failed to register a consumer")
	forever := make(chan bool)
	go func() {
		for d := range msgs {
			var target Target
			err := json.Unmarshal(d.Body, &target)
			failOnError(err, "Failed to deal with json")
			status := scan(target.IP, target.Port)
			fmt.Printf("[+] Scanning %s:%s\n", target.IP, target.Port)
			response := fmt.Sprintf("%s --> %s is %s", target.IP, target.Port, status)
			err = ch.Publish(
				"",        // exchange
				d.ReplyTo, // routing key
				false,     // mandatory,
				false,     // immediate
				amqp.Publishing{
					ContentType:   "application/json",
					CorrelationId: d.CorrelationId,
					Body:          []byte(response),
				},
			)
			failOnError(err, "Failed to publish a message")
			d.Ack(false)
		}
	}()
	log.Printf("[*] Awaiting RPC requests")
	<-forever
}
